#!/usr/bin/env python
# coding: utf-8

import sys
import shutil
import json
from pathlib import Path
from loguru import logger
from ruamel.yaml import YAML
from labelme2yolo_detect import Coco2Yolo
import dataset_splite
from app_config import APP_CFG
from create_train_yaml import CreateTrainYaml
from train_model import modelTrain

from runtimer import runtimer_create
from model_export import modelExport
# from predict_rddhw import TRddhwPredict
from bhclasses_define import BHClasses

APP_FILE = Path(sys.argv[0])
APP_ROOT = APP_FILE.parent
APP_NAME = APP_FILE.stem
LOG_FILE = APP_ROOT / "log" / f"{APP_NAME}.log"
LOG_FILE_ERR = APP_ROOT / "log" / f"{APP_NAME}-err.log"
LOG_FILE_HANDLER = None
LOG_FILE_ERR_HANDLER = None
MODEL_ROOT = Path("/aaby/netdisk/mydata/ai-data/model/rddhw")  # 最终模型输出的目录



def _logInit() -> None:
    global LOG_FILE_HANDLER, LOG_FILE_ERR_HANDLER

    # logger.add(sys.stderr, format="{time} {level} {message}", level="INFO")
    # time = datetime.datetime.today()
    # f'the time is {time:%Y-%m-%d (%a) %H:%M:%S}'   # datetime时间格式
    LOG_FILE_HANDLER = logger.add(LOG_FILE, rotation="500 MB", encoding='utf-8')  
    LOG_FILE_ERR_HANDLER = logger.add(LOG_FILE_ERR, rotation="500 MB", encoding='utf-8', level='ERROR')  
    # logger.add(f"{APP_ROOT}/log/{APP_NAME}-{time:%Y%m%d}.log", rotation="500 MB", encoding='utf-8', retention="10 days")  
    # logger.add(f"{APP_ROOT}/log/{APP_NAME}-err-{time:%Y%m%d}.log", rotation="500 MB", encoding='utf-8', level='ERROR', retention="30 days")  

    # t = time.strftime("%Y%m%d") 
    # logger.add(f"{baseDir}/log-{t}.log", rotation="500 MB", encoding='utf-8', retention="10 days")  
    # logger.add(f"{baseDir}/err-{t}.log", rotation="500 MB", encoding='utf-8', level='ERROR', retention="30 days")  




def _train(train_file: Path)->str:
    proj_name = modelTrain(APP_CFG.model_cfg, train_file, MODEL_ROOT)    
    return proj_name

def main():
    logger.info(f"训练流程开始：{APP_NAME}")
    # === 配置 =========================
    # classes_file = APP_ROOT / "classes-rddhw-labelme.json"  # 类别定义文件，和 labelme 标注的标签一致，类别的ID必须从0开始
    # classes_file_labelme = APP_ROOT / "classes-rddhw-labelme.json"   # labelme 标注格式用的 classes 定义文件
    # 初始化配置
    APP_CFG.load(APP_ROOT / "app-debug.yaml")
    img_root_train = Path(APP_CFG.pipe_cfg.img_root_train)
    data_root_predict = Path(APP_CFG.pipe_cfg.data_root_predict)
    train_file = img_root_train.parent / "train.yaml"

    # ==== 配置参数检查 =================
    # if not classes_file.exists():
    #     print(f"[ERR] 类别定义文件不存在：{classes_file}")
    #     return 

    # APP_CFG.pipe_cfg.datainit = False

    # ==== run =========================
    if APP_CFG.pipe_cfg.datainit:
        # 1. 将 labelme 的 json 标注文件转换为 yolo 格式的 txt 文件
        logger.info("1. 将 labelme 的 json 标注文件转换为 yolo 格式的 txt 文件：")
        Coco2Yolo.l2y_dir(img_root_train, BHClasses.dict_name_id)
        logger.info(".")

        # 2. 将 json 标注文件中嵌入的图片删除
        # logger.info("2. 将 labelme 标注文件 .json 中嵌入的图片删除：")
        # l2y.json_delimgs(img_root_train)
        # logger.info(".")


        # 3. 自动划分数据集到训练集、验证集和测试集
        logger.info("3. 自动划分数据集到训练集、验证集和测试集：")
        dataset_splite.datasetSplit(img_root_train)
        logger.info(".")

        # 4. 创建 train.yaml 文件
        logger.info("4. 创建 train.yaml 文件：")
        CreateTrainYaml(img_root_train.parent, BHClasses.dict_id_name)
        logger.info(".")
        pass
    else:
        logger.info("跳过数据准备阶段")    
    
    # 5. 开始训练 [计时]
    logger.info("5. 开始训练：")
    with runtimer_create("训练"):
        proj_name = _train(train_file)
    logger.info(".")

    # 6. 导出模型 [计时]
    model_file = MODEL_ROOT / proj_name / "weights/best.pt"
    logger.info(f"6. 导出模型：{str(model_file)}")
    with runtimer_create("导出模型"):
        modelExport(model_file, APP_CFG.model_cfg.imgsz)
    logger.info(".")

    # 7. 用生成的 engine 模型进行预测 [计时] (不推荐此时预测，需要人工确认置信度参数)
    # model_engine = model_file.with_suffix(".engine")
    # logger.info("7. 用生成的 engine 模型进行预测")
    # with runtimer_create("预测自动标注数据") as timer:
    #     file_sum, bh_sum = predictRddhw(model_engine, data_root_predict, classes_file_labelme, APP_CFG.model_cfg.imgsz, APP_CFG.pipe_cfg.conf)
    # logger.info(f"识别完成，文件总数：{file_sum}, 识别速度：{float(file_sum)/timer.elapsed * 60:.2f} 张/分, 识别病害文件总数：{bh_sum}")

    # end. 结束，将日志文件移动到模型目录下
    logger.info("训练流程结束")
    logger.remove(LOG_FILE_HANDLER)
    logger.remove(LOG_FILE_ERR_HANDLER)
    log_dir = MODEL_ROOT / proj_name / "log"
    log_dir.mkdir(parents=True, exist_ok=True)
    shutil.move(LOG_FILE, log_dir/LOG_FILE.name)
    shutil.move(LOG_FILE_ERR, log_dir/LOG_FILE_ERR.name)
    pass

if __name__=="__main__":
    _logInit()
    main()

